op {
  name: "BoostedTreesTrainingPredict"
  input_arg {
    name: "tree_ensemble_handle"
    type: DT_RESOURCE
  }
  input_arg {
    name: "cached_tree_ids"
    type: DT_INT32
  }
  input_arg {
    name: "cached_node_ids"
    type: DT_INT32
  }
  input_arg {
    name: "bucketized_features"
    type: DT_INT32
    number_attr: "num_bucketized_features"
  }
  output_arg {
    name: "partial_logits"
    type: DT_FLOAT
  }
  output_arg {
    name: "tree_ids"
    type: DT_INT32
  }
  output_arg {
    name: "node_ids"
    type: DT_INT32
  }
  attr {
    name: "num_bucketized_features"
    type: "int"
    has_minimum: true
    minimum: 1
  }
  attr {
    name: "logits_dimension"
    type: "int"
  }
  is_stateful: true
}
